package aes

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"fmt"
	"io"
)

// 补码
func PKCS7Padding(ciphertext []byte, blockSize int) []byte {
	padding := blockSize - len(ciphertext)%blockSize
	padText := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(ciphertext, padText...)
}

// 去码
func PKCS7UnPadding(originData []byte) []byte {
	length := len(originData)
	unPadding := int(originData[length-1])
	return originData[:(length - unPadding)]
}

// CBC模式下的AES加密
func CBCEncrypt(data []byte, key []byte) ([]byte, error) {
	data = PKCS7Padding(data, aes.BlockSize)
	block, _ := aes.NewCipher(key)
	out := make([]byte, aes.BlockSize+len(data))
	iv := out[:aes.BlockSize]
	_, _ = io.ReadFull(rand.Reader, iv)
	mode := cipher.NewCBCEncrypter(block, iv)
	mode.CryptBlocks(out[aes.BlockSize:], data)
	return out, nil
}

// CBC模式下的AES解密
func CBCDecrypt(data []byte, key []byte) ([]byte, error) {
	block, _ := aes.NewCipher(key)
	iv := data[:aes.BlockSize]
	data = data[aes.BlockSize:]
	if len(data)%aes.BlockSize != 0 {
		return nil, fmt.Errorf("data is not a multiple of the block size")
	}
	out := make([]byte, len(data))
	mode := cipher.NewCBCDecrypter(block, iv)
	mode.CryptBlocks(out, data)
	out = PKCS7UnPadding(out)
	return out, nil
}

// CFB模式下的AES加密
func CFBEncrypt(plainText []byte, key []byte) []byte {
	// TODO 异常处理
	// 分组密钥
	block, _ := aes.NewCipher(key)
	// 创建用于存放加密后密文的数组
	cipherText := make([]byte, aes.BlockSize+len(plainText))
	// 设置内存空间可读
	iv := cipherText[:aes.BlockSize]
	// TODO 异常处理
	_, _ = io.ReadFull(rand.Reader, iv)
	// 设置加密模式, 会返回一个流
	stream := cipher.NewCFBEncrypter(block, iv)
	// 加密利用cipherText[aes.BlockSize]与plainText进行异或运算
	stream.XORKeyStream(cipherText[aes.BlockSize:], plainText)
	// 返回密文数据
	return cipherText
}

// CFB模式下的AES解密
func CFBDecryptCFB(cipherText []byte, key []byte) []byte {
	// TODO 异常处理
	// 分组密钥
	block, _ := aes.NewCipher(key)
	iv := cipherText[:aes.BlockSize]
	cipherText = cipherText[aes.BlockSize:]
	// 设置解密方式
	stream := cipher.NewCFBDecrypter(block, iv)
	stream.XORKeyStream(cipherText, cipherText)
	return cipherText
}

// OFB模式下的AES加密
func OFBEncrypt(plainText []byte, key []byte) []byte {
	plainText = PKCS7Padding(plainText, aes.BlockSize)
	block, _ := aes.NewCipher(key)
	crypt := make([]byte, aes.BlockSize+len(plainText))
	iv := crypt[:aes.BlockSize]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		panic(err.Error())
	}
	stream := cipher.NewOFB(block, iv)
	stream.XORKeyStream(crypt[aes.BlockSize:], plainText)
	return crypt
}

// OFB模式下的AES解密
func OFBDecrypt(crypt []byte, key []byte) []byte {
	block, _ := aes.NewCipher(key)
	iv := crypt[:aes.BlockSize]
	crypt = crypt[aes.BlockSize:]

	if len(crypt)%aes.BlockSize != 0 {
		panic("data is not a multiple of the block size")
	}
	plainText := make([]byte, len(crypt))
	blockMode := cipher.NewOFB(block, iv)
	blockMode.XORKeyStream(plainText, crypt)
	return plainText
}

// CTR模式下的AES加密
func CTREncrypt(plainText []byte, key []byte) []byte {
	plainText = PKCS7Padding(plainText, aes.BlockSize)
	block, _ := aes.NewCipher(key)
	crypt := make([]byte, aes.BlockSize+len(plainText))
	iv := crypt[:aes.BlockSize]
	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
		panic("")
	}
	stream := cipher.NewCTR(block, iv)
	stream.XORKeyStream(crypt[aes.BlockSize:], plainText)
	return crypt
}

// CTR模式下的AES解密
func CTRDecrypt(crypt []byte, key []byte) []byte {
	block, _ := aes.NewCipher(key)
	iv := crypt[:aes.BlockSize]
	crypt = crypt[aes.BlockSize:]
	if len(crypt)%aes.BlockSize != 0 {
		panic("")
	}
	plainText := make([]byte, len(crypt))
	blockMode := cipher.NewCTR(block, iv)
	blockMode.XORKeyStream(plainText, crypt)
	plainText = PKCS7UnPadding(plainText)
	return plainText
}

// ECB模式下的AES加密
func ECBEncrypt(plainText []byte, key []byte) []byte {
	// 分组密钥
	block, _ := aes.NewCipher(key)
	// 获得密钥块的长度
	blockSize := block.BlockSize()
	// 补码
	plainText = PKCS7Padding(plainText, blockSize)
	// 创建加密后的密文数组
	crypt := make([]byte, len(plainText))
	//
	for bs, be := 0, blockSize; bs < len(plainText); bs, be = bs+blockSize, be+blockSize {
		block.Encrypt(crypt[bs:be], plainText[bs:be])
	}
	return crypt
}

// ECB模式下的AES解密
func ECBDecrypt(crypt []byte, key []byte) []byte {
	block, _ := aes.NewCipher(key)
	decrypted := make([]byte, len(crypt))
	size := block.BlockSize()

	for bs, be := 0, size; bs < len(crypt); bs, be = bs+size, be+size {
		block.Decrypt(decrypted[bs:be], crypt[bs:be])
	}

	return PKCS7UnPadding(decrypted)
}

// GCM模式下的AES加密
func GCMEncrypt(key []byte, data []byte) ([]byte, error) {
	block, _ := aes.NewCipher(key)
	blockMode, _ := cipher.NewGCM(block)

	nonce := make([]byte, blockMode.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		panic(err.Error())
	}

	out := blockMode.Seal(nonce, nonce, data, nil)
	return out, nil
}

// GCM模式下的AES解密
func GCMDecrypt(key []byte, data []byte) ([]byte, error) {
	block, _ := aes.NewCipher(key)
	blockMode, _ := cipher.NewGCM(block)

	nonce, ciphertext := data[:blockMode.NonceSize()], data[blockMode.NonceSize():]
	out, _ := blockMode.Open(nil, nonce, ciphertext, nil)

	return out, nil
}
